package OceanModel;

public class GriddedVariable {
	  // Class representing model data on a regular grid
	  // Provides methods for data interpolation and finding variable derivatives
	
	  double[][][] data;
	  int nx, ny, nz;
	  double dx, dy;
	  int halo;
	  boolean reentrantX, reentrantY;
	  char gridType;
	  Parameters params;
	  
	  GriddedVariable(Parameters params, char initGridType) {
		  // Sets up the model grid using information in params
          // Grid is a staggered Arakawa C-grid, which looks like this:
		  // 2 q v q v q
		  // 1 u h u h u
		  // 1 q v q v q
		  // 0 u h u h u
		  // 0 q v q v q		
		  //   0 0 1 1 2
	    this.params = params;
	    gridType = initGridType;
	    dx = params.dx;
	    dy = params.dy;
	    halo = params.halo;
	    reentrantX = params.reentrantX;
	    reentrantY = params.reentrantY;

	    nx = params.nx + 2*params.halo;
	    ny = params.ny + 2*params.halo;
	    nz = params.nz;

	    switch(gridType) {
	      case 'h':
	        break;
	      case 'u':
	        nx = nx+1;
	        break;
	      case 'v':
	        ny = ny+1;
	        break;
	      case 'q':
	        nx = nx+1;
	        ny = ny+1;
	        break;
	    }

	    data = new double[nz][ny][nx];
	  } // GriddedVariable

	  //----------------------

	  public double[][][] getData() {
		// Returns the model domain data, discarding the halo points
	    double[][][] result;
	    
	    result = new double[nz][ny-2*halo][nx-2*halo];
	    for (int k=0; k < nz; k++) {      
	      for (int j=halo; j < ny-halo; j++) {      
	        for (int i=halo; i < nx-halo; i++) {
	          result[k][j-halo][i-halo] = data[k][j][i];
	        }
	      }
	    }
	    
	    return result;
	  } // getData

	  public double getData(int k, int j, int i) {
	    // Returns the model domain data, discarding the halo points, at a single point
	    return data[k][j+halo][i+halo];
	  } // getData

	  //----------------------

	  GriddedVariable onGrid(char new_grid) {
		// Interpolates GriddedVariable data onto another grid
	    GriddedVariable result;
	    
	    result = new GriddedVariable(params, new_grid);

	    switch(gridType) {
	      case 'h':
	        switch(new_grid) {
	          case 'h':
	            result = this;
	            break;
	          case 'u':
	            result = hToU();
	            break;
	          case 'v':
	            result = hToV();
	            break;
	          case 'q':
	            result = hToQ();
	            break;
	        }
	        break;
	      case 'u':
	        switch(new_grid) {
	          case 'h':
	            result = uToH();
	            break;
	          case 'u':
	            result = this;
	            break;
	          case 'v':
	            result = uToV();
	            break;
	          case 'q':
	            result = uToQ();
	            break;
	        }
	        break;
	      case 'v':
	        switch(new_grid) {
	          case 'h':
	            result = vToH();
	            break;
	          case 'u':
	            result = vToU();
	            break;
	          case 'v':
	            result = this;
	            break;
	          case 'q':
	            result = vToQ();
	            break;
	        }
	        break;
	      case 'q':
	        switch(new_grid) {
	          case 'h':
	            result = qToH();
	            break;
	          case 'u':
	            result = qToU();
	            break;
	          case 'v':
	            result = qToV();
	            break;
	          case 'q':
	            result = this;
	            break;
	        }
	        break;
	    }

	    result.swapHalo();
	    
	    return result;
	  } // onGrid

	  // The next functions interpolate model variables from one grid type to another	  
	  GriddedVariable hToU() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'u');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny; j++) {
	        for (int i=1; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j][i-1] + data[k][j][i])/2.;
	        }
	      }
	    }
	    return result;
	  } // hToU

	  GriddedVariable hToV() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'v');
	    for (int k=0; k < nz; k++) {
	      for (int j=1; j < ny; j++) {
	        for (int i=0; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j-1][i] + data[k][j][i])/2.;
	        }
	      }
	    }
	    return result;
	  } // hToV

	  GriddedVariable hToQ() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'q');
	    for (int k=0; k < nz; k++) {
	      for (int j=1; j < ny; j++) {
	        for (int i=1; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j-1][i-1] + data[k][j-1][i] + data[k][j][i-1] + data[k][j][i])/4.;
	        }
	      }
	    }
	    return result;
	  } // hToQ

	  GriddedVariable uToH() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'h');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny; j++) {
	        for (int i=0; i < nx-1; i++) {
	          result.data[k][j][i] = (data[k][j][i] + data[k][j][i+1])/2.;
	        }
	      }
	    }
	    return result;
	  } // uToH

	  GriddedVariable uToV() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'v');
	    for (int k=0; k < nz; k++) {
	      for (int j=1; j < ny; j++) {
	        for (int i=0; i < nx-1; i++) {
	          result.data[k][j][i] = (data[k][j-1][i] + data[k][j][i] + data[k][j-1][i+1] + data[k][j][i+1])/4.;
	        }
	      }
	    }
	    return result;
	  } // uToV

	  GriddedVariable uToQ() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'q');
	    for (int k=0; k < nz; k++) {
	      for (int j=1; j < ny; j++) {
	        for (int i=0; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j][i] + data[k][j+1][i])/2.;
	        }
	      }
	    }
	    return result;
	  } // uToQ


	  GriddedVariable vToH() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'h');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny-1; j++) {
	        for (int i=0; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j][i] + data[k][j+1][i])/2.;
	        }
	      }
	    }
	    return result;
	  } // vToH

	  GriddedVariable vToU() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'u');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny-1; j++) {
	        for (int i=1; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j][i-1] + data[k][j+1][i-1] + data[k][j][i] + data[k][j+1][i])/4.;
	        }
	      }
	    }
	    return result;
	  } // vToU

	  GriddedVariable vToQ() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'q');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny; j++) {
	        for (int i=1; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j][i] + data[k][j][i+1])/2.;
	        }
	      }
	    }
	    return result;
	  } // vToQ

	  GriddedVariable qToH() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'h');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny-1; j++) {
	        for (int i=0; i < nx-1; i++) {
	          result.data[k][j][i] = (data[k][j][i] + data[k][j+1][i] + data[k][j][i+1] + data[k][j+1][i+1])/4.;
	        }
	      }
	    }
	    return result;
	  } // qToH

	  GriddedVariable qToU() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'u');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny-1; j++) {
	        for (int i=0; i < nx; i++) {
	          result.data[k][j][i] = (data[k][j][i] + data[k][j+1][i])/2.;
	        }
	      }
	    }
	    return result;
	  } // qToU

	  GriddedVariable qToV() {
	    GriddedVariable result;
	    result = new GriddedVariable(params, 'v');
	    for (int k=0; k < nz; k++) {
	      for (int j=0; j < ny; j++) {
	        for (int i=0; i < nx-1; i++) {
	          result.data[k][j][i] = (data[k][j][i] + data[k][j][i+1])/2.;
	        }
	      }
	    }
	    return result;
	  } // qToV


	  //----------------------

	  GriddedVariable d_dx() {
		  // Returns the derivative of the variable in the x direction
		  // The derivative is a first order difference
	    GriddedVariable result;
	    char grid_result = gridType;
	    
	    switch(gridType) {
	      case 'h':
	        grid_result = 'u';
	        break;
	      case 'v':
	        grid_result = 'q';
	        break;
	      case 'u':
	        grid_result = 'h';
	        break;
	      case 'q':
	        grid_result = 'v';
	        break;
	    }

	    result = new GriddedVariable(params, grid_result);

	    switch(gridType) {
	      case 'h':
	      case 'v':
	        for (int k=0; k < nz; k++) {
	          for (int j=0; j < ny; j++) {
	            for (int i=1; i < nx; i++) {
	              result.data[k][j][i] = (data[k][j][i] - data[k][j][i-1])/dx;
	            }
	          }
	        }          
	        break;
	      case 'u':
	      case 'q':
	        for (int k=0; k < nz; k++) {
	          for (int j=0; j < ny; j++) {
	            for (int i=0; i < nx-1; i++) {
	              result.data[k][j][i] = (data[k][j][i+1] - data[k][j][i])/dx;
	            }
	          }
	        }          
	        break;
	    }

	    result.swapHalo();

	    return result;
	  } // d_dx

	  //----------------------

	  GriddedVariable d_dy() {
		  // Returns the derivative of the variable in the y direction
		  // The derivative is a first order difference
	    GriddedVariable result;
	    char grid_result = gridType;

	    switch(gridType) {
	      case 'h':
	        grid_result = 'v';
	        break;
	      case 'v':
	        grid_result = 'h';
	        break;
	      case 'u':
	        grid_result = 'q';
	        break;
	      case 'q':
	        grid_result = 'u';
	        break;
	    }

	    result = new GriddedVariable(params, grid_result);

	    switch(gridType) {
	      case 'h':
	      case 'u':
	        for (int k=0; k < nz; k++) {
	          for (int j=1; j < ny; j++) {
	            for (int i=0; i < nx; i++) {
	              result.data[k][j][i] = (data[k][j][i] - data[k][j-1][i])/dy;
	            }
	          }
	        }          
	        break;
	      case 'v':
	      case 'q':
	        for (int k=0; k < nz; k++) {
	          for (int j=0; j < ny-1; j++) {
	            for (int i=0; i < nx; i++) {
	              result.data[k][j][i] = (data[k][j+1][i] - data[k][j][i])/dy;
	            }
	          }
	        }          
	        break;
	    }

	    result.swapHalo();

	    return result;
	  } // d_dy
	  
	  //----------------------

	  GriddedVariable d2_dx2() {
		  // Returns the second derivative of the variable in the x direction
		  // The derivative is a first order difference
	    GriddedVariable result;
	    
	    result = new GriddedVariable(params, gridType);

	    for (int k=0; k < nz; k++) {
	      for (int j=1; j < ny-1; j++) {
	        for (int i=1; i < nx-1; i++) {
	          result.data[k][j][i] = (data[k][j][i+1] + data[k][j][i-1] - 2.*data[k][j][i])/dx/dx;
	        }
	      }
	    }          

	    result.swapHalo();
	    
	    return result;
	  } // d2_dx2

	  //----------------------

	  GriddedVariable d2_dy2() {
		  // Returns the second derivative of the variable in the y direction
		  // The derivative is a first order difference
	    GriddedVariable result;
	    
	    result = new GriddedVariable(params, gridType);

	    for (int k=0; k < nz; k++) {
	      for (int j=1; j < ny-1; j++) {
	        for (int i=1; i < nx-1; i++) {
	          result.data[k][j][i] = (data[k][j+1][i] + data[k][j-1][i] - 2.*data[k][j][i])/dy/dy;
	        }
	      }
	    }          

	    result.swapHalo();
	    
	    return result;
	  } // d2_dy2

	  //----------------------

	  void swapHalo() {
		  // If the domain is reentrant, sets halo points to the variable values on the 
		  // other side of the domain to simulate periodic boundary conditions.
		  // If the domain is not reentrant, ensures the halo points are always 0		  
		  // This should be run after any changes happen to the variable.
	    if (reentrantX) {
	      for (int k=0; k < nz; k++) {
	        for (int i=0; i < halo; i++) {
	          for (int j=0; j < ny; j++) {
	            data[k][j][i] = data[k][j][nx - 2*halo + i];
	            data[k][j][nx - halo + i] = data[k][j][i + halo];
	          }
	        }
	      }
	    }
	    else {
	      for (int k=0; k < nz; k++) {
	        for (int i=0; i < halo; i++) {
	          for (int j=0; j < ny; j++) {
	            data[k][j][i] = 0.;
	            data[k][j][nx - halo + i] = 0.;
	          }
	        }
	      }
	    }

	    if (reentrantY) {
	      for (int k=0; k < nz; k++) {
	        for (int i=0; i < nx; i++) {
	          for (int j=0; j < halo; j++) {
	            data[k][j][i] = data[k][ny - 2*halo + j][i];
	            data[k][ny - halo + j][i] = data[k][j + halo][i];
	          }
	        }
	      }
	    }
	    else {
	      for (int k=0; k < nz; k++) {
	        for (int i=0; i < nx; i++) {
	          for (int j=0; j < halo; j++) {
	            data[k][j][i] = 0.;
	            data[k][ny - halo + j][i] = 0.;
	          }
	        }
	      }      
	    }
	    
	  } // swapHalo

}
